import torch
from torchvision import transforms
from config import cfg

t = transforms.Compose([
    # H W C --> C H W 且把值归一化为 0-1
    transforms.ToTensor(),
])


def bbox_iou(box, boxes):
    box_area = (box[2] - box[0]) * (box[3] - box[1])
    boxes_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    l_x = torch.maximum(box[0], boxes[:, 0])
    l_y = torch.maximum(box[1], boxes[:, 1])
    r_x = torch.minimum(box[2], boxes[:, 2])
    r_y = torch.minimum(box[3], boxes[:, 3])
    w = torch.maximum(r_x - l_x, torch.tensor(0))
    h = torch.maximum(r_y - l_y, torch.tensor(0))
    inter_area = w * h
    iou_val = inter_area / (box_area + boxes_area - inter_area)
    return iou_val


def nms(detect_boxes, threshold=0.5):
    """
    :param detect_boxes: 侦测输出的框的信息 [[conf, tx, ty, tw, th, cls], ...]
    :param threshold: 阈值
    :return: 筛选后的侦测框
    流程分析
    1. 模型输出的框，按置信度排序
    2. 置信度最高的，作为当前类别最优的框 max_conf_box = detect_boxes[0]
    3. 剩余的框 detect_boxes[1:] 和当前最优框 max_conf_box 计算IOU 获取 iou_val
    4. 和给定阈值(超参数)作比较 iou_idx = iou_val < thresh
    5. detect_boxes[1:][iou_idx] 则为保留的框
    """
    # 保留最优框信息
    best_boxes = []
    # 1. 模型输出的框，按置信度排序
    idx = torch.argsort(detect_boxes[:, 0], descending=True)
    detect_boxes = detect_boxes[idx]
    while detect_boxes.size(0) > 0:
        # 2. 置信度最高的，作为当前类别最优的框
        max_conf_box = detect_boxes[0]
        best_boxes.append(max_conf_box)
        # 3. 剩余的框 detect_boxes[1:] 和当前最优框 max_conf_box 计算IOU
        detect_boxes = detect_boxes[1:]
        iou_val = bbox_iou(max_conf_box[1:5], detect_boxes[:, 1:5])
        # 4. 和给定阈值(超参数)作比较保留小于阈值的对应框
        detect_boxes = detect_boxes[iou_val < threshold]
    return best_boxes
